#! /usr/bin/env python
# coding=utf-8


# 在一个ogs文件中找到趋同位点

# 参数 
pep_cut_ali = r'pep_cut_ali.fa'


import argparse
import sys

from ete3 import PhyloTree
from Bio import SeqIO

import itertools
# 读取group
from scipy.stats import fisher_exact


parser = argparse.ArgumentParser(
    description="""从 aaml祖先重建的结果中读取 氨基酸的趋同替换 需要准备一个文件 叫 con_group 里面是各个组 每个组 的物种用空格空开 一行一个组
    允许有包含关系的组
    用法:\n
    conver03_find_conver_single.py  -i OGS_for_second_step -c con_group -o convergent_site\n
    由大天才于2021年11月12日创建于浙江农业大学""")


parser.add_argument('-i',
                help='必须给定,运行完aaml的ogs目录')


parser.add_argument('-o',
                help='必须给定,输出的路径')

parser.add_argument('-c',
                help='必须给定,group 文件')


#outfile = r'convergent_site'
#con_group  = r'con_group' # 趋同进化的group


args = parser.parse_args()

if not args.i or not args.o or not args.c:
    parser.print_help()
    sys.exit()



infile = args.i
#outpath = r'OGS_for_second_step'
outfile = args.o
con_group = args.c




group_dic = {}


with open(con_group) as fila:
    n = 1
    for i in fila:
        k = i.strip().split(' ')
        if k!=[''] and k !=[]:
            group_dic['group_'+str(n)] = k
            n += 1
            # print(k)
#print(group_dic)
# 需要在文件里包含rst 为aaml 祖先重建的结果
# 读取带数字的进化tree




# group 包含关系
#print(group_dic)

def run_one_file(infile):
    #print(infile)
    tree = ''
    with open(infile+'/rst') as fila:
        t =  0
        for i in fila:
            if t ==1:
                tree = i.strip()
                break
            if i.find('''tree with node labels for Rod Page's TreeView''') != -1:
                t = 1
    if tree.strip() =='':
        print('tree error at '+infile)
        return []
    ptree = PhyloTree(tree)

    # 获取group对应的节点


    # 这里的中间节点名称就是 support
    group_node_dic = {}
     
    def find_group_node(t):
        #print(t)
        if t.is_leaf():
            node_id = t.name.split('_')[0]
        else:
            node_id = str(int(t.support))

        leaf_name_lista = [i.split('_')[1] for i in  t.get_leaf_names()]
        
        for i in group_dic:
            

            l = group_dic[i]
            

            q = 1
            for ii in leaf_name_lista:

                if ii not in l:
                    q = 0
                    break

            if q == 1:

                if i not in group_node_dic:
                    #print(i,l,node_id)
                    group_node_dic[i] = node_id
                
        for i in t.get_children():
            
            find_group_node(i)
        #for i in group_node_dic:
    



    find_group_node(ptree)

    #print(group_node_dic)
    #node_group_dic = {}
    # 对每个group 选找转变
   


    node_group_dic = {k:v for v,k in group_node_dic.items()}
    #print(node_group_dic)

    group_change_dic = {i:{} for i in group_node_dic}

    #print(group_change_dic)

    with open(infile+'/rst') as fila:
        t =  0
        node_id =None
        for i in fila:
            i = i.strip()
            # if i =="":
            #     continue
            if i.find("Check root of tree for directions of change.") !=-1:
                t = 1


            if t==1:
                if i.find('Branch') ==0:
                    node_id = i.split('Branch')[1].split('..')[1].strip().split(' ')[0]
                    nnn = 0
                else:
                    if node_id in node_group_dic and nnn < 2:

                        if len(i.strip().split(' '))< 2:
                            nnn += 1
                            continue

                        group_id = node_group_dic[node_id]
                        #print(group_id,i,node_id)
                        #print(group_id)
                        item = i.split('->')[-1].strip()
                        if item == '-':
                            continue
                        site = int(i.split(' ')[0])-1
                        raw_site = i.split(' ')[1]
                        target_site = i.split(' ')[4]
                        group_change_dic[group_id][site] = (site,raw_site,target_site)




            if i.find("List of extant and reconstructed sequences") !=-1:
                break

    # 对group进行组合
    #print(group_change_dic)

    group_lista = list(group_change_dic)

    #print(node_group_dic)



    # 列出所有组对的排列组合

    comb_lista = []
    for i in range(2,len(group_lista)):
        k = itertools.combinations(group_lista , i)
        for i in k:
            comb_lista.append(i)



    # 读取所有seq
    seq_dic = {str(i.name):str(i.seq) for i in SeqIO.parse(infile+'/'+pep_cut_ali,'fasta')}

    spe_lista = [str(i.name) for i in SeqIO.parse(infile+'/'+pep_cut_ali,'fasta')]
    con_spe_lista = []
    for i in group_dic:
        con_spe_lista += group_dic[i]
    #non_con_spe_lista = [for i in str]



    # 根据group 判断 哪些是convergent 位点


    target_site = []

    #print(group_change_dic)


    #print(comb_lista)
    # 遍历所有的组合 将得到的结果存入 target_site 中

    for lista in comb_lista:
        #final_site = []


        # com_site 是所有变化位点的并集 用于一个一个查找
        com_site = []
        for j in lista:
    #         # 先找并 再找交集
            if j in group_change_dic:
                k = group_change_dic[j].keys()
                for i in k: 
                    if i not in com_site:
                        com_site.append(i)
        
      
        for i in com_site:

            t = 1

            # 改变在两个类群中必须都存在
            for j in lista:

                if i not in group_change_dic[j]:
            
                    t = 0
            
                    break


                #group_change_dic[j][2]
            #print(group_change_dic)
            # 改变的结果在两个类群中必须一致
            if t ==1:
                #print(lista)
                aa = group_change_dic[lista[0]][i][2]
                for j in lista[1:]:
                    #print(j)
                    if group_change_dic[j][i][2] != aa:
                        t =0
                        break

            # 趋同变化的氨基酸 不允许是 -
            if t ==1 and aa!='-':

                # 要求对象物种都是
                for j in lista:
                    target_spe_lista = group_dic[j]
                    for k in target_spe_lista:
                        if k not in seq_dic:
                            continue
                        if seq_dic[k][i] != aa:
                            t = 0
                            break
            else:
                continue

            # 在非趋同物种 使用fisher extract 检验 在非趋同物种中 该氨基酸的出现频率是否比总体的低
            if t ==1:
                con_site = []
                not_con_site = []
                for j in seq_dic:
                    #print(j)
                    if j not in con_spe_lista:
                        not_con_site.append(seq_dic[j][i]==aa)
                    else:
                        con_site.append(seq_dic[j][i]==aa)
                #print(total_site)
                #print(not_con_site)
                con_site_true = con_site.count(True)
                con_site_false = con_site.count(False)
                not_con_site_true = not_con_site.count(True)
                not_con_site_false = not_con_site.count(False)
                
                #oddsratio, pvalue = fisher_exact([[con_site_true , not_con_site_true], [con_site_false, not_con_site_false]],'greater')
                pvalue =  not_con_site_true / (not_con_site_true+not_con_site_false)

            if t ==1:
               
                target_site.append((i,lista,aa,pvalue))

    
    # 去处重复

    a = 0
    b = 0
    #print(target_site)
    while a < len(target_site):
        b = a+1
        n = 0
        while b < len(target_site):
            if target_site[a][0] == target_site[b][0] and target_site[b][2] ==target_site[a][2]:
                t = 0
                for g in target_site[a][1]:
                    if g not in target_site[b][1]:
                        t = 1
                        break
                if t == 0:
                    target_site.remove(target_site[a])
                    n = 1
                    break
            b += 1
        if n == 0:
            a += 1

    # 去除包含关系的group

    new_target_site =[]
    for j in target_site:
        t = 0
        for k in j[1]:
            for d in j[1]:
                if k!=d:
                    k_spe = set(group_dic[k]) & set(spe_lista)
                    d_spe = set(group_dic[d]) & set(spe_lista)
                    #print(k_spe,d_spe)
                    if k_spe.issubset(d_spe) or k_spe.issuperset(d_spe):
                        t = 1
                        break
        if t == 0:
            new_target_site.append(j)
    target_site = new_target_site


    write_lines = []

    for k in target_site:


        write_line = ''
        group_line = ''
        group_l = []
        target_spe = None
        context = ''

        for j in k[1]:
            #print(j)
            group_l.append(j+': '+' '.join([i for i in group_dic[j] if i in seq_dic]))
            if target_spe==None:
                target_spe = [i for i in group_dic[j] if i in seq_dic][0]
        group_line =' | '.join(group_l)
        context = seq_dic[target_spe][max([0,k[0]-10]):min([k[0]+10,len(seq_dic[target_spe])])]



        write_line = infile.split('/')[-1] +'\t'+str(k[0])+'\t'+ group_line + '\t'+k[2]+'\t'+str(k[3])+'\t'+context+'\t'+str(len(group_l))
        # 还需要提取地一个group物种中上下游5bp的序列 用于锁定 
        write_lines.append(write_line)
    return write_lines



    # 检查目标是否具有一致性

#print(group_dic)


        
            
#         else:
#             continue
        

# print(group_change_dic)
import os

with open(outfile ,'w') as fila:

    n = 1
    for f in os.listdir(infile):

        target_file = infile+'/'+f

        w = run_one_file(target_file)

        n += 1
        for i in w:
            fila.write(i+'\n')

#run_one_file('OGS_for_second_step/OG0002710')